from rapidocr_download_url import ocr_and_extract_links
import os
from config import DEFAULT_PAGE_RANGE, DEFAULT_CONTAIN_STRING, DEFAULT_FIND_STOP

def process_pdf_list(input_file, output_file, page_range=DEFAULT_PAGE_RANGE, 
                    contain_string=DEFAULT_CONTAIN_STRING, find_stop=DEFAULT_FIND_STOP):
    """
    处理PDF链接列表文件，提取链接并分类存储
    
    Args:
        input_file (str): 包含PDF链接的输入文件路径
        output_file (str): 输出结果文件路径
        page_range (str): 页面范围，默认为"1-2"
        contain_string (list): 包含的关键字列表，默认为["https", ".com", ".io"]
        find_stop (bool): 是否在找到链接后停止识别后续页面，默认为True
    """
    # 检查输入文件是否存在
    if not os.path.exists(input_file):
        print(f"错误: 输入文件 {input_file} 不存在")
        return
    
    # 读取PDF链接列表
    with open(input_file, 'r', encoding='utf-8') as f:
        pdf_urls = [line.strip() for line in f.readlines() if line.strip()]
    
    print(f"读取到 {len(pdf_urls)} 个PDF链接")
    
    # 准备结果列表
    results = []
    
    # 处理每个PDF链接
    for i, pdf_url in enumerate(pdf_urls, 1):
        print(f"\n处理第 {i}/{len(pdf_urls)} 个PDF: {pdf_url}")
        
        try:
            # 使用OCR处理PDF并提取链接，使用传入的参数
            result = ocr_and_extract_links(
                file_url=pdf_url,
                page_range=page_range,  # 使用传入的页面范围
                contain_string=contain_string,  # 使用传入的关键字列表
                find_stop=find_stop  # 使用传入的find_stop参数
            )
            
            # 获取提取到的链接
            extracted_links = result.get('extracted_links', [])
            
            if extracted_links:
                print(f"  从 {pdf_url} 提取到 {len(extracted_links)} 个链接")
                
                # 分类链接
                github_links = []
                other_links = []
                
                for link in extracted_links:
                    if "github.com" in link:
                        github_links.append(link)
                    else:
                        other_links.append(link)
                
                # 添加到结果中
                results.append({
                    'pdf_url': pdf_url,
                    'github_links': github_links,
                    'other_links': other_links
                })
            else:
                print(f"  从 {pdf_url} 未提取到链接")
                results.append({
                    'pdf_url': pdf_url,
                    'github_links': [],
                    'other_links': []
                })
                
        except Exception as e:
            print(f"  处理 {pdf_url} 时出错: {e}")
            results.append({
                'pdf_url': pdf_url,
                'github_links': [],
                'other_links': []
            })
    
    # 将结果写入输出文件
    write_results_to_file(results, output_file)
    print(f"\n结果已保存到 {output_file}")

def write_results_to_file(results, output_file):
    """
    将结果写入文件
    
    Args:
        results (list): 处理结果列表
        output_file (str): 输出文件路径
    """
    with open(output_file, 'w', encoding='utf-8') as f:
        # 写入表头
        f.write("PDF链接\t包含github.com的链接\t其他链接\n")
        
        # 写入数据
        for result in results:
            pdf_url = result['pdf_url']
            github_links = result['github_links']
            other_links = result['other_links']
            
            # 如果某类链接为空，写入空字符串
            github_str = '; '.join(github_links) if github_links else ''
            other_str = '; '.join(other_links) if other_links else ''
            
            f.write(f"{pdf_url}\t{github_str}\t{other_str}\n")

def main():
    # 输入文件路径
    input_file = "input_files/pdf_url_list.txt"
    # 输出文件路径
    output_file = "output_links_classification.txt"
    
    print("开始处理PDF链接列表...")
    process_pdf_list(input_file, output_file)
    print("处理完成!")

if __name__ == "__main__":
    main()